import runpy
import sys
import argparse
from rich.progress import Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
from rich.console import Console


console = Console()

# parse CLI for optional monthly GIF generation
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('--no-monthly-gif', action='store_true', help='跳过在 try.py 中生成月度 GIF（默认生成）')
parser.add_argument('--monthly-scope', choices=['county', 'city', 'all'], default='county', help='月度 GIF 的范围（默认 county）')
parser.add_argument('--monthly-n-workers', type=int, default=8, help='生成月度 GIF 时使用的并行 worker 数（默认 8）')
parser.add_argument('--monthly-out', default=None, help='月度 GIF 输出路径，默认根据 scope 自动生成')
ARGS, _REMAIN = parser.parse_known_args()


def main() -> int:
    """Run try_plot.py while showing a rich spinner + elapsed time.

    Returns an exit code (0 on success)."""
    with Progress(
        SpinnerColumn(),
        TextColumn("[bold blue]运行 try_plot.py…"),
        TimeElapsedColumn(),
        transient=True,
    ) as progress:
        task_id = progress.add_task("run", total=None)
        try:
            runpy.run_path("try_plot.py", run_name="__main__")
            # optionally generate monthly GIF after try_plot finishes
            if not ARGS.no_monthly_gif:
                # build argv for scripts.generate_time_series_gif.main
                scope = ARGS.monthly_scope
                n_workers = ARGS.monthly_n_workers
                out = ARGS.monthly_out
                # decide out filename if not provided
                if out is None:
                    if scope == 'county':
                        out = 'figures/云南_边境县_月度灯光_全量_from_try.gif'
                    elif scope == 'city':
                        out = 'figures/云南_边境城市_月度灯光_全量_from_try.gif'
                    else:
                        out = 'figures/云南_全量_月度灯光_全量_from_try.gif'

                # assemble args for the gif script
                gif_args = [
                    '--parquet', 'data/云南乡镇月度灯光亮度.parquet',
                    '--population-parquet', 'data/边境城市七普乡镇.parquet',
                    '--value-col', '日度夜间灯光亮度',
                    '--out', out,
                    '--n-workers', str(n_workers),
                    '--verbose'
                ]
                if scope == 'county':
                    gif_args.append('--only-border-counties')
                elif scope == 'city':
                    gif_args.append('--only-border-cities')

                try:
                    # import and call the main function of the gif script
                    from scripts.generate_time_series_gif import main as gif_main
                    console.log(f'开始在 try.py 中调用月度 GIF 生成，scope={scope}, n_workers={n_workers}, out={out}')
                    gif_main(gif_args)
                    console.log(f'月度 GIF 生成完成: {out}')
                except Exception as e:
                    console.log(f'生成月度 GIF 失败：{e}')
            return 0
        finally:
            # ensure spinner stops promptly
            try:
                progress.stop_task(task_id)
            except Exception:
                pass


if __name__ == "__main__":
    from rich.console import Console
    from rich.panel import Panel
    from rich.text import Text
    import time as _time
    # custom start panel for try.py to visually distinguish it from other scripts
    start_time = _time.time()
    start_time_str = _time.strftime('%Y-%m-%d %H:%M:%S')
    Console().print(Panel(Text(f"⏱  运行 try_plot 开始 {start_time_str}", style="bold white on dark_blue"), expand=False))
    try:
        code = main()
        sys.exit(code if isinstance(code, int) else 0)
    except SystemExit as e:
        code = e.code if isinstance(e.code, int) else 0
        sys.exit(code)
    except Exception as e:
        Console().log(f'运行 try.py 出错：{e}')
        import traceback
        Console().log(traceback.format_exc())
        sys.exit(1)
    finally:
        end_time = _time.time()
        end_time_str = _time.strftime('%Y-%m-%d %H:%M:%S')
        duration = end_time - start_time
        body = Text()
        body.append(f"⌛️  运行 try_plot 结束 {end_time_str}\n", style="bold white")
        body.append(f"总耗时：{int(duration)} 秒（{duration/60:.2f} 分钟）", style="yellow")
        Console().print(Panel(body, style="on dark_blue", expand=False))
